from celery import shared_task,Task

from dvadmin.knowledge.models import KnowledgeRecord,KnowledgeRecordSegment
from dvadmin.embedding.embedding import Bge
from dvadmin.db.milvus import MilvusUtil

@shared_task
def deal_file(rid):
    bge = Bge()
    milvus = MilvusUtil()
    knowledgeRecordSegments = KnowledgeRecordSegment.objects.filter(rid=rid)
    print('开始处理文件')
                
    chunkList=[]
    data_vector=[]
    
    for knowledgeRecordSegment in knowledgeRecordSegments:
        kid = knowledgeRecordSegment.kid
        chunkList.append(knowledgeRecordSegment.content)
        data_vector.append({
            "id": knowledgeRecordSegment.id,
            "kid": kid,
            "file_id": knowledgeRecordSegment.rid,
            "file_valid": 1,
            "chunk_valid": 1,   
            "context": knowledgeRecordSegment.content
        })
    
    # 向量化
    embeddings = bge.embed_documents(chunkList)
    print('向量化')
    
    row = 0
    for embedding in embeddings:            
        data_vector[row]["embedding"] = embedding
        row += 1
    
    tablename = "knowledge" + str(kid)
    # 保存到向量数据库
    milvus.upset(tablename,data_vector)
    print('保存到向量数据库')
    
    knowledgeRecord = KnowledgeRecord.objects.get(id=rid)
    knowledgeRecord.status=1
    knowledgeRecord.save()